from diffusers import AutoPipelineForImage2Image
from datetime import datetime
import torch
import gc

class Image2ImageModel:
  def __init__(self, model_id: str):
    pipeline = AutoPipelineForImage2Image.from_pretrained(model_id, local_files_only=True,  torch_dtype=torch.float16,  use_safetensors=True)
    pipeline = pipeline.to("cuda")
    # pipe.enable_model_cpu_offload()
    self.pipeline = pipeline
  
  def generate(self, prompt: str, negative_prompt: str, init_image):
    image = self.pipeline(
      prompt=prompt,
      negative_prompt=negative_prompt,
      image=init_image,
      guidance_scale=3.0,
      width=1024,
      height=1024,
      safety_checker=False,
      strength=0.5,
      num_inference_steps=50
    ).images[0]
    # image.save("output_image.png")
    return image
  
  def save_image(self, prompt: str, negative_prompt: str, init_image, output_dir: str = ""):
    image = self.generate(prompt, negative_prompt, init_image)
    ts = datetime.now().strftime("%Y%m%d%H%M%S")
    filename = f"{output_dir}{ts}.png"
    image.save(filename)
    
  def mult_gen(self, num_images_per_prompt:int, prompt: str, negative_prompt: str, init_image):
    steps = 75
    scale = 7
    seed = torch.randint(0, 1000000, (1,)).item()
    generator = torch.Generator(device="cuda").manual_seed(seed)
    images = self.pipeline(prompt, negative_prompt=negative_prompt, image=init_image, width=1024, height=1024, num_inference_steps=steps,
    guidance_scale=scale, num_images_per_prompt=num_images_per_prompt, generator=generator).images
    return images
  
  def close(self):
    del self.pipeline
    gc.collect()
    torch.cuda.empty_cache()
    

__Image2ImageModelInstance__: Image2ImageModel = None

def image2image(model_id: str, prompt: str, negative_prompt: str, init_image):
  global __Image2ImageModelInstance__
  if __Image2ImageModelInstance__ is None:
    print("create module instance..")
    __Image2ImageModelInstance__ = Image2ImageModel(model_id)
  return __Image2ImageModelInstance__.generate(prompt, negative_prompt, init_image)

def image2imageN(model_id: str, num_images_per_prompt: int, prompt: str, negative_prompt: str, init_image):
  model = Image2ImageModel(model_id)
  # image = model.generate(prompt)
  images = model.mult_gen(num_images_per_prompt, prompt, negative_prompt, init_image)
  model.close()
  print(f"create module instance {num_images_per_prompt} times..")
  return images